package main

import (
	"bufio"
	"fmt"
	"os"
	"strings"

	"generateTOTP/storage"
	"generateTOTP/totp"
)

func main() {
	reader := bufio.NewReader(os.Stdin)

	for {
		fmt.Println("\n===== Google Authenticator TOTP 程序 =====")
		fmt.Println("a. 生成一个密钥")
		fmt.Println("b. 校验动态码")
		fmt.Println("c. 删除密钥")
		fmt.Println("d. 生成动态码") // 添加新选项
		fmt.Println("q. 退出程序")
		fmt.Print("请选择操作 (a/b/c/d/q): ")

		choice, _ := reader.ReadString('\n')
		choice = strings.TrimSpace(choice)

		switch choice {
		case "a":
			generateKey(reader)
		case "b":
			verifyCode(reader)
		case "c":
			deleteKey(reader)
		case "d":
			generateCode(reader) // 添加新的case处理
		case "q":
			fmt.Println("程序已退出")
			return
		default:
			fmt.Println("无效的选择，请重新输入")
		}
	}
}

func generateKey(reader *bufio.Reader) {
	fmt.Print("请输入账号: ")
	account, _ := reader.ReadString('\n')
	account = strings.TrimSpace(account)

	if account == "" {
		fmt.Println("账号不能为空")
		return
	}

	// 生成密钥
	secret, err := totp.GenerateSecret()
	if err != nil {
		fmt.Printf("生成密钥失败: %v\n", err)
		return
	}

	// 保存密钥
	err = storage.SaveSecret(account, secret)
	if err != nil {
		fmt.Printf("保存密钥失败: %v\n", err)
		return
	}

	// 生成二维码链接
	qrURL := totp.GenerateQRCodeURL("TOTP Demo", account, secret)

	fmt.Println("\n===== 密钥信息 =====")
	fmt.Printf("账号: %s\n", account)
	fmt.Printf("密钥: %s\n", secret)
	fmt.Printf("二维码链接: %s\n", qrURL)
	fmt.Println("\n请使用 Google Authenticator 扫描上面的二维码链接，或手动添加密钥")
}

func verifyCode(reader *bufio.Reader) {
	fmt.Print("请输入账号: ")
	account, _ := reader.ReadString('\n')
	account = strings.TrimSpace(account)

	if account == "" {
		fmt.Println("账号不能为空")
		return
	}

	// 获取密钥
	secret, err := storage.GetSecret(account)
	if err != nil {
		fmt.Printf("获取密钥失败: %v\n", err)
		return
	}

	fmt.Print("请输入 Google Authenticator 生成的动态验证码: ")
	code, _ := reader.ReadString('\n')
	code = strings.TrimSpace(code)

	// 验证动态码
	valid, err := totp.ValidateCode(secret, code)
	if err != nil {
		fmt.Printf("验证失败: %v\n", err)
		return
	}

	if valid {
		fmt.Println("验证成功！动态码匹配")
	} else {
		fmt.Println("验证失败！动态码不匹配")
	}
}

// 在文件末尾添加删除密钥的函数
func deleteKey(reader *bufio.Reader) {
	fmt.Print("请输入要删除的账号: ")
	account, _ := reader.ReadString('\n')
	account = strings.TrimSpace(account)

	if account == "" {
		fmt.Println("账号不能为空")
		return
	}

	// 删除密钥
	err := storage.DeleteSecret(account)
	if err != nil {
		fmt.Printf("删除密钥失败: %v\n", err)
		return
	}

	fmt.Printf("账号 %s 的密钥已成功删除\n", account)
}

// 添加生成动态码的函数
func generateCode(reader *bufio.Reader) {
	fmt.Print("请输入账号: ")
	account, _ := reader.ReadString('\n')
	account = strings.TrimSpace(account)

	if account == "" {
		fmt.Println("账号不能为空")
		return
	}

	// 获取密钥
	secret, err := storage.GetSecret(account)
	if err != nil {
		fmt.Printf("获取密钥失败: %v\n", err)
		return
	}

	// 生成当前的动态码
	code, err := totp.GenerateCode(secret)
	if err != nil {
		fmt.Printf("生成动态码失败: %v\n", err)
		return
	}

	fmt.Println("\n===== 动态码信息 =====")
	fmt.Printf("账号: %s\n", account)
	fmt.Printf("当前动态码: %s\n", code)
}
